# doubao_tts_client.py
# 豆包大模型TTS WebSocket客户端

import asyncio
import websockets
import uuid
import json
import gzip
from settings import settings

class DoubaoTTSClient:
    def __init__(self, appid, token, cluster="volcano_tts", voice_type="zh_male_aojiaobazong_emo_v2_mars_bigtts"):
        self.appid = appid
        self.token = token
        self.cluster = cluster
        self.voice_type = voice_type
        self.host = "openspeech.bytedance.com"
        self.api_url = f"wss://{self.host}/api/v1/tts/ws_binary"
        
        # 默认头部：version(4) + header_size(4) + message_type(4) + flags(4) + serialization(4) + compression(4) + reserved(8)
        # 0b0001 0001 0001 0000 0001 0001 00000000 = 0x11 0x10 0x11 0x00
        self.default_header = bytearray(b'\x11\x10\x11\x00')

    async def stream_text_to_audio(self, text):
        """
        流式合成文本为音频，返回音频数据块
        """
        header = {"Authorization": f"Bearer; {self.token}"}
        
        # 构建请求JSON
        request_json = {
            "app": {
                "appid": self.appid,
                "token": "access_token",
                "cluster": self.cluster
            },
            "user": {
                "uid": str(uuid.uuid4())
            },
            "audio": {
                "voice_type": self.voice_type,
                "encoding": "pcm",
                "speed_ratio": 1.0,
                "volume_ratio": 1.0,
                "pitch_ratio": 1.0,
                "rate": 8000,
            },
            "request": {
                "reqid": str(uuid.uuid4()),
                "text": text,
                "text_type": "plain",
                "operation": "submit"
            }
        }
        
        print(f"[豆包TTS] 请求参数: {json.dumps(request_json, indent=2)}")
        
        # 构建二进制请求
        payload_bytes = str.encode(json.dumps(request_json))
        payload_bytes = gzip.compress(payload_bytes)
        full_client_request = bytearray(self.default_header)
        full_client_request.extend((len(payload_bytes)).to_bytes(4, 'big'))  # payload size(4 bytes)
        full_client_request.extend(payload_bytes)  # payload
        
        print(f"[豆包TTS] 发送请求，长度: {len(full_client_request)} 字节")
        
        try:
            async with websockets.connect(self.api_url, additional_headers=header, ping_interval=None) as ws:
                await ws.send(full_client_request)
                
                while True:
                    res = await ws.recv()
                    done, audio_data = self._parse_response(res)
                    
                    if audio_data:
                        yield audio_data
                    
                    if done:
                        print(f"[豆包TTS] 合成完成")
                        break
                        
        except Exception as e:
            print(f"[豆包TTS] 连接异常: {e}")
            raise

    def _parse_response(self, res):
        """
        解析豆包TTS响应
        返回: (是否完成, 音频数据)
        """
        try:
            # 解析头部
            protocol_version = res[0] >> 4
            header_size = res[0] & 0x0f
            message_type = res[1] >> 4
            message_type_specific_flags = res[1] & 0x0f
            serialization_method = res[2] >> 4
            message_compression = res[2] & 0x0f
            reserved = res[3]
            
            header_extensions = res[4:header_size*4]
            payload = res[header_size*4:]
            
            print(f"[豆包TTS] 响应解析: message_type={message_type:#x}, flags={message_type_specific_flags:#x}")
            
            # 处理音频响应 (0xb = 11)
            if message_type == 0xb:  # audio-only server response
                if message_type_specific_flags == 0:  # no sequence number as ACK
                    print(f"[豆包TTS] 收到ACK响应")
                    return False, None
                else:
                    sequence_number = int.from_bytes(payload[:4], "big", signed=True)
                    payload_size = int.from_bytes(payload[4:8], "big", signed=False)
                    audio_data = payload[8:]
                    
                    print(f"[豆包TTS] 音频段: sequence={sequence_number}, size={payload_size}")
                    
                    # 负数表示合成完毕
                    if sequence_number < 0:
                        return True, audio_data
                    else:
                        return False, audio_data
                        
            # 处理错误响应 (0xf = 15)
            elif message_type == 0xf:
                code = int.from_bytes(payload[:4], "big", signed=False)
                msg_size = int.from_bytes(payload[4:8], "big", signed=False)
                error_msg = payload[8:]
                
                if message_compression == 1:
                    error_msg = gzip.decompress(error_msg)
                error_msg = str(error_msg, "utf-8")
                
                print(f"[豆包TTS] 错误: code={code}, message={error_msg}")
                raise Exception(f"豆包TTS错误: code={code}, message={error_msg}")
                
            # 处理前端响应 (0xc = 12)
            elif message_type == 0xc:
                msg_size = int.from_bytes(payload[:4], "big", signed=False)
                payload = payload[4:]
                
                if message_compression == 1:
                    payload = gzip.decompress(payload)
                    
                print(f"[豆包TTS] 前端消息: {payload}")
                return False, None
                
            else:
                print(f"[豆包TTS] 未知消息类型: {message_type:#x}")
                return True, None
                
        except Exception as e:
            print(f"[豆包TTS] 响应解析异常: {e}")
            raise 